#include "yuvviewer.h"
#include <QMenu>
#include <QAction>
#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QPushButton>
#include <QLabel>

// 顶点数据：位置和纹理坐标
GLfloat vertices[] = {
    // 位置          // 纹理坐标
    -1.0f, -1.0f, 0.0f,  0.0f, 1.0f,
    1.0f, -1.0f, 0.0f,  1.0f, 1.0f,
    1.0f,  1.0f, 0.0f,  1.0f, 0.0f,
    -1.0f,  1.0f, 0.0f,  0.0f, 0.0f
};

GLuint indices[] = {
    0, 1, 2,
    2, 3, 0
};

// 顶点着色器源码
const char *vertexShaderSource = R"(
#version 330 core
layout (location = 0) in vec3 aPos;
layout (location = 1) in vec2 aTexCoord;

out vec2 TexCoord;

void main()
{
    gl_Position = vec4(aPos, 1.0);
    TexCoord = aTexCoord;
})";


// 片段着色器源码 - YUV转RGB
const char *fragmentShaderSource = R"(
#version 330 core
in vec2 TexCoord;
out vec4 FragColor;

uniform sampler2D texY;
uniform sampler2D texU;
uniform sampler2D texV;

// BT.601转换矩阵  标准范围
const mat3 yuv2rgb = mat3(
    1.164383,  1.164383, 1.164383,
    0.000000, -0.391762, 2.017232,
    1.596027, -0.812968, 0.000000
);

//全范围，目标测试文件是YUV420全范围，用这个效果最好
const mat3 yuv2rgb2 = mat3(
    1,     1,       1,
    0,     -0.3441, 1.772,
    1.402, -0.7141, 0
);

//测试
const mat3 yuv2rgb3 = mat3(
    1,       1,         1,
    0,       -0.39465,  2.03211,
    1.13983, -0.58060,  0
);

void main()
{
    float y = texture(texY, TexCoord).r;
    float u = texture(texU, TexCoord).r - 0.5;
    float v = texture(texV, TexCoord).r - 0.5;

    vec3 yuv = vec3(y, u, v);

    vec3 rgb = yuv2rgb2 * yuv;
    FragColor = vec4(rgb, 1.0);
}
)";




YUVViewer::YUVViewer(QWidget *parent)
    : QOpenGLWidget(parent), shaderProgram(nullptr),
    textureY(nullptr), textureU(nullptr), textureV(nullptr),
    zoomLevel(1.0f), rotationX(0.0f), rotationY(0.0f),
    imageWidth(640), imageHeight(480)
{
    setFocusPolicy(Qt::StrongFocus);

    // 创建控制面板
    QWidget *controlPanel = new QWidget(this);
    QVBoxLayout *mainLayout = new QVBoxLayout(this);
    QHBoxLayout *buttonLayout = new QHBoxLayout();

    QPushButton *openButton = new QPushButton("打开YUV文件", this);
    connect(openButton, &QPushButton::clicked, this, &YUVViewer::openYUVFile);

    buttonLayout->addWidget(openButton);
    buttonLayout->addStretch();

    controlPanel->setLayout(buttonLayout);
    mainLayout->addWidget(controlPanel);
    mainLayout->addWidget(this);

    VBO = 0;
    VAO = 0;
    EBO = 0;
}

YUVViewer::~YUVViewer()
{
    makeCurrent();

    if(shaderProgram)
    {
        shaderProgram->release();
        delete shaderProgram;
    }

    if(textureY)
    {
        textureY->destroy();
        delete textureY;
    }
    if(textureU)
    {
        textureU->destroy();
        delete textureU;
    }
    if(textureV)
    {
        textureV->destroy();
        delete textureV;
    }

    if(VAO) glDeleteVertexArrays(1, &VAO);
    if(VBO) glDeleteBuffers(1, &VBO);
    if(EBO) glDeleteBuffers(1, &EBO);

    doneCurrent();
}

void YUVViewer::initializeGL()
{
    initializeOpenGLFunctions();
    //glClearColor(0.1f, 0.1f, 0.1f, 1.0f);
    glClearColor(1.0f, 1.0f, 1.0f, 1.0f);

    setupVertexData();
    setupShaders();
}

void YUVViewer::resizeGL(int w, int h)
{
    glViewport(0, 0, w, h);
}

void YUVViewer::paintGL()
{
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

    shaderProgram->bind();

    // 绑定YUV纹理
    if(textureY && textureY->isCreated())
    {
        glActiveTexture(GL_TEXTURE0);
        textureY->bind();
        shaderProgram->setUniformValue("texY", 0);
    }

    if(textureU && textureU->isCreated())
    {
        glActiveTexture(GL_TEXTURE1);
        textureU->bind();
        shaderProgram->setUniformValue("texU", 1);
    }

    if(textureV && textureV->isCreated())
    {
        glActiveTexture(GL_TEXTURE2);
        textureV->bind();
        shaderProgram->setUniformValue("texV", 2);
    }

    glBindVertexArray(VAO);
    glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, 0);
    glBindVertexArray(0);

    shaderProgram->release();
}

void YUVViewer::setupVertexData()
{
    glGenVertexArrays(1, &VAO);
    glGenBuffers(1, &VBO);
    glGenBuffers(1, &EBO);

    glBindVertexArray(VAO);

    glBindBuffer(GL_ARRAY_BUFFER, VBO);
    glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(indices), indices, GL_STATIC_DRAW);

    // 位置属性
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 5 * sizeof(GLfloat), (GLvoid*)0);
    glEnableVertexAttribArray(0);

    // 纹理坐标属性
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 5 * sizeof(GLfloat), (GLvoid*)(3 * sizeof(GLfloat)));
    glEnableVertexAttribArray(1);

    glBindVertexArray(0);
}

void YUVViewer::setupShaders()
{
    shaderProgram = new QOpenGLShaderProgram(this);

    // 编译顶点着色器
    if(!shaderProgram->addShaderFromSourceCode(QOpenGLShader::Vertex, vertexShaderSource))
    {
        qDebug() << "Vertex shader compilation failed:" << shaderProgram->log();
        return;
    }

    // 编译片段着色器
    if(!shaderProgram->addShaderFromSourceCode(QOpenGLShader::Fragment, fragmentShaderSource))
    {
        qDebug() << "Fragment shader compilation failed:" << shaderProgram->log();
        return;
    }

    // 链接程序
    if(!shaderProgram->link())
    {
        qDebug() << "Shader program linking failed:" << shaderProgram->log();
        return;
    }
}

void YUVViewer::openYUVFile()
{
    QString fileName = QFileDialog::getOpenFileName(this,
                                                    tr("打开YUV文件"), "", tr("YUV文件 (*.yuv)"));

    if(!fileName.isEmpty())
    {
        QFile file(fileName);
        if(!file.open(QIODevice::ReadOnly))
        {
            QMessageBox::information(this, tr("YUV查看器"),
                                     tr("无法打开文件 %1.").arg(fileName));
            return;
        }

        yuvData = file.readAll();
        file.close();

        if(yuvData.isEmpty())
        {
            QMessageBox::information(this, tr("YUV查看器"),
                                     tr("文件为空 %1.").arg(fileName));
            return;
        }

        // 假设YUV420P格式，imageWidth(640), imageHeight(480)
        imageWidth = 640;
        imageHeight = 480;
        int frameSize = imageWidth * imageHeight * 3 / 2;

        if(yuvData.size() < frameSize)
        {
            QMessageBox::information(this, tr("YUV查看器"),
                                     tr("文件大小不匹配，期望大小：%1，实际大小：%2")
                                         .arg(frameSize).arg(yuvData.size()));
            return;
        }

        makeCurrent();

        // 清理现有纹理
        if(textureY)
        {
            textureY->destroy();
            delete textureY;
            textureY = nullptr;
        }
        if(textureU)
        {
            textureU->destroy();
            delete textureU;
            textureU = nullptr;
        }
        if(textureV)
        {
            textureV->destroy();
            delete textureV;
            textureV = nullptr;
        }

        //使用纯QOpenGLTexture接口（推荐）
        // 创建Y纹理
        textureY = new QOpenGLTexture(QOpenGLTexture::Target2D);
        textureY->create();
        textureY->setSize(imageWidth, imageHeight);
        textureY->setFormat(QOpenGLTexture::R8_UNorm);  // 使用R8格式
        textureY->allocateStorage();
        textureY->setData(QOpenGLTexture::Red, QOpenGLTexture::UInt8, yuvData.constData());
        textureY->setMinificationFilter(QOpenGLTexture::Linear);
        textureY->setMagnificationFilter(QOpenGLTexture::Linear);
        textureY->setWrapMode(QOpenGLTexture::ClampToEdge);

        // 创建U纹理
        textureU = new QOpenGLTexture(QOpenGLTexture::Target2D);
        textureU->create();
        textureU->setSize(imageWidth/2, imageHeight/2);
        textureU->setFormat(QOpenGLTexture::R8_UNorm);
        textureU->allocateStorage();
        textureU->setData(QOpenGLTexture::Red, QOpenGLTexture::UInt8,
                          yuvData.constData() + imageWidth * imageHeight);
        textureU->setMinificationFilter(QOpenGLTexture::Linear);
        textureU->setMagnificationFilter(QOpenGLTexture::Linear);
        textureU->setWrapMode(QOpenGLTexture::ClampToEdge);

        // 创建V纹理
        textureV = new QOpenGLTexture(QOpenGLTexture::Target2D);
        textureV->create();
        textureV->setSize(imageWidth/2, imageHeight/2);
        textureV->setFormat(QOpenGLTexture::R8_UNorm);
        textureV->allocateStorage();
        textureV->setData(QOpenGLTexture::Red, QOpenGLTexture::UInt8,
                          yuvData.constData() + imageWidth * imageHeight + (imageWidth/2) * (imageHeight/2));
        textureV->setMinificationFilter(QOpenGLTexture::Linear);
        textureV->setMagnificationFilter(QOpenGLTexture::Linear);
        textureV->setWrapMode(QOpenGLTexture::ClampToEdge);

        update();
    }
}

void YUVViewer::mousePressEvent(QMouseEvent *event)
{
    lastMousePos = event->pos();
}

void YUVViewer::mouseMoveEvent(QMouseEvent *event)
{
    if(event->buttons() & Qt::LeftButton)
    {
        int dx = event->x() - lastMousePos.x();
        int dy = event->y() - lastMousePos.y();

        rotationY += dx * 0.5f;
        rotationX += dy * 0.5f;

        lastMousePos = event->pos();
        update();
    }
}

void YUVViewer::wheelEvent(QWheelEvent *event)
{
    QPoint numDegrees = event->angleDelta() / 8;

    if(!numDegrees.isNull())
    {
        if(numDegrees.y() > 0)
        {
            zoomLevel *= 1.1f;
        }
        else
        {
            zoomLevel /= 1.1f;
        }

        zoomLevel = qMax(0.1f, qMin(5.0f, zoomLevel));
        update();
    }
}
